package system

import (
	"encoding/json"
	"fmt"
	"log"
	"sync"
	"time"

	"github.com/asynkron/protoactor-go/actor"
)

// StateManager is an actor that manages state
type StateManager struct {
	// stateStore is the underlying state store
	stateStore StateStore
	// snapshots are the state snapshots
	snapshots map[string]*StateSnapshot
	// subscribers are the actors subscribed to state changes
	subscribers map[string][]*actor.PID
	// mu is a mutex for protecting the state
	mu sync.RWMutex
	// logger is the state manager logger
	logger *log.Logger
	// checkpointInterval is the interval between automatic checkpoints
	checkpointInterval time.Duration
	// lastCheckpoint is when the last checkpoint was created
	lastCheckpoint time.Time
}

// StateStore is an interface for state storage
type StateStore interface {
	// Get gets a state value
	Get(key string) ([]byte, error)
	// Put puts a state value
	Put(key string, value []byte) error
	// Delete deletes a state value
	Delete(key string) error
	// CreateCheckpoint creates a checkpoint
	CreateCheckpoint(id string) error
	// RestoreFromCheckpoint restores from a checkpoint
	RestoreFromCheckpoint(id string) error
	// ListCheckpoints lists all checkpoints
	ListCheckpoints() ([]string, error)
}

// StateSnapshot represents a snapshot of state
type StateSnapshot struct {
	// ID is the snapshot ID
	ID string
	// CreatedAt is when the snapshot was created
	CreatedAt time.Time
	// States is the state data
	States map[string][]byte
}

// NewStateManager creates a new state manager
func NewStateManager() actor.Actor {
	return &StateManager{
		stateStore:         NewInMemoryStateStore(),
		snapshots:          make(map[string]*StateSnapshot),
		subscribers:        make(map[string][]*actor.PID),
		logger:             log.New(log.Writer(), "[StateManager] ", log.LstdFlags),
		checkpointInterval: 5 * time.Minute,
		lastCheckpoint:     time.Now(),
	}
}

// Receive handles messages for the state manager
func (s *StateManager) Receive(context actor.Context) {
	switch msg := context.Message().(type) {
	case *actor.Started:
		s.handleStarted(context)
	case *actor.Stopping:
		s.handleStopping(context)
	case *SaveStateCommand:
		s.handleSaveState(context, msg)
	case *LoadStateCommand:
		s.handleLoadState(context, msg)
	case *CreateCheckpointCommand:
		s.handleCreateCheckpoint(context, msg)
	case *RestoreFromCheckpointCommand:
		s.handleRestoreFromCheckpoint(context, msg)
	case *ListCheckpointsCommand:
		s.handleListCheckpoints(context)
	case *SubscribeToStateChangesCommand:
		s.handleSubscribeToStateChanges(context, msg)
	case *UnsubscribeFromStateChangesCommand:
		s.handleUnsubscribeFromStateChanges(context, msg)
	case *CheckpointTimerMessage:
		s.handleCheckpointTimer(context)
	default:
		s.logger.Printf("Received unknown message: %T", msg)
	}
}

// handleStarted handles the Started message
func (s *StateManager) handleStarted(context actor.Context) {
	s.logger.Println("State manager started")

	// Start checkpoint timer
	context.Send(context.Self(), &CheckpointTimerMessage{})
}

// handleStopping handles the Stopping message
func (s *StateManager) handleStopping(context actor.Context) {
	s.logger.Println("State manager stopping")

	// Create final checkpoint
	checkpointID := fmt.Sprintf("final-%d", time.Now().Unix())
	err := s.stateStore.CreateCheckpoint(checkpointID)
	if err != nil {
		s.logger.Printf("Failed to create final checkpoint: %v", err)
	} else {
		s.logger.Printf("Created final checkpoint: %s", checkpointID)
	}
}

// handleSaveState handles the SaveStateCommand message
func (s *StateManager) handleSaveState(context actor.Context, msg *SaveStateCommand) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Convert state to JSON
	data, err := json.Marshal(msg.State)
	if err != nil {
		context.Respond(&SaveStateResponse{
			Success: false,
			Error:   fmt.Errorf("failed to marshal state: %w", err),
		})
		return
	}

	// Save state
	err = s.stateStore.Put(msg.Key, data)
	if err != nil {
		context.Respond(&SaveStateResponse{
			Success: false,
			Error:   fmt.Errorf("failed to save state: %w", err),
		})
		return
	}

	// Respond with success
	context.Respond(&SaveStateResponse{
		Success: true,
	})

	// Notify subscribers
	s.notifyStateChanged(context, msg.Key, msg.State)

	s.logger.Printf("Saved state for key: %s", msg.Key)
}

// handleLoadState handles the LoadStateCommand message
func (s *StateManager) handleLoadState(context actor.Context, msg *LoadStateCommand) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	// Load state
	data, err := s.stateStore.Get(msg.Key)
	if err != nil {
		context.Respond(&LoadStateResponse{
			Success: false,
			Error:   fmt.Errorf("failed to load state: %w", err),
		})
		return
	}

	// No data found
	if data == nil {
		context.Respond(&LoadStateResponse{
			Success: true,
			State:   nil,
		})
		return
	}

	// Unmarshal state
	var state interface{}
	err = json.Unmarshal(data, &state)
	if err != nil {
		context.Respond(&LoadStateResponse{
			Success: false,
			Error:   fmt.Errorf("failed to unmarshal state: %w", err),
		})
		return
	}

	// Respond with state
	context.Respond(&LoadStateResponse{
		Success: true,
		State:   state,
	})

	s.logger.Printf("Loaded state for key: %s", msg.Key)
}

// handleCreateCheckpoint handles the CreateCheckpointCommand message
func (s *StateManager) handleCreateCheckpoint(context actor.Context, msg *CreateCheckpointCommand) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Create checkpoint
	err := s.stateStore.CreateCheckpoint(msg.ID)
	if err != nil {
		context.Respond(&CreateCheckpointResponse{
			Success: false,
			Error:   fmt.Errorf("failed to create checkpoint: %w", err),
		})
		return
	}

	// Update last checkpoint time
	s.lastCheckpoint = time.Now()

	// Respond with success
	context.Respond(&CreateCheckpointResponse{
		Success: true,
	})

	s.logger.Printf("Created checkpoint: %s", msg.ID)
}

// handleRestoreFromCheckpoint handles the RestoreFromCheckpointCommand message
func (s *StateManager) handleRestoreFromCheckpoint(context actor.Context, msg *RestoreFromCheckpointCommand) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Restore from checkpoint
	err := s.stateStore.RestoreFromCheckpoint(msg.ID)
	if err != nil {
		context.Respond(&RestoreFromCheckpointResponse{
			Success: false,
			Error:   fmt.Errorf("failed to restore from checkpoint: %w", err),
		})
		return
	}

	// Respond with success
	context.Respond(&RestoreFromCheckpointResponse{
		Success: true,
	})

	s.logger.Printf("Restored from checkpoint: %s", msg.ID)
}

// handleListCheckpoints handles the ListCheckpointsCommand message
func (s *StateManager) handleListCheckpoints(context actor.Context) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	// List checkpoints
	checkpoints, err := s.stateStore.ListCheckpoints()
	if err != nil {
		context.Respond(&ListCheckpointsResponse{
			Success: false,
			Error:   fmt.Errorf("failed to list checkpoints: %w", err),
		})
		return
	}

	// Respond with checkpoints
	context.Respond(&ListCheckpointsResponse{
		Success:     true,
		Checkpoints: checkpoints,
	})
}

// handleSubscribeToStateChanges handles the SubscribeToStateChangesCommand message
func (s *StateManager) handleSubscribeToStateChanges(context actor.Context, msg *SubscribeToStateChangesCommand) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Add subscriber
	s.subscribers[msg.Key] = append(s.subscribers[msg.Key], context.Sender())

	// Respond with success
	context.Respond(&SubscribeToStateChangesResponse{
		Success: true,
	})

	s.logger.Printf("Added subscriber for key %s: %s", msg.Key, context.Sender().String())
}

// handleUnsubscribeFromStateChanges handles the UnsubscribeFromStateChangesCommand message
func (s *StateManager) handleUnsubscribeFromStateChanges(context actor.Context, msg *UnsubscribeFromStateChangesCommand) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Find and remove subscriber
	subscribers := s.subscribers[msg.Key]
	for i, pid := range subscribers {
		if pid.String() == context.Sender().String() {
			s.subscribers[msg.Key] = append(subscribers[:i], subscribers[i+1:]...)
			break
		}
	}

	// Respond with success
	context.Respond(&UnsubscribeFromStateChangesResponse{
		Success: true,
	})

	s.logger.Printf("Removed subscriber for key %s: %s", msg.Key, context.Sender().String())
}

// handleCheckpointTimer handles the CheckpointTimerMessage message
func (s *StateManager) handleCheckpointTimer(context actor.Context) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Check if it's time for a checkpoint
	now := time.Now()
	if now.Sub(s.lastCheckpoint) >= s.checkpointInterval {
		// Create checkpoint
		checkpointID := fmt.Sprintf("auto-%d", now.Unix())
		err := s.stateStore.CreateCheckpoint(checkpointID)
		if err != nil {
			s.logger.Printf("Failed to create automatic checkpoint: %v", err)
		} else {
			s.logger.Printf("Created automatic checkpoint: %s", checkpointID)
			s.lastCheckpoint = now
		}
	}

	// Schedule next checkpoint check
	context.Send(context.Self(), &CheckpointTimerMessage{})
}

// notifyStateChanged notifies subscribers of a state change
func (s *StateManager) notifyStateChanged(context actor.Context, key string, newState interface{}) {
	// Get subscribers for this key
	subscribers := s.subscribers[key]
	if len(subscribers) == 0 {
		return
	}

	// Create event
	event := &StateChangedEvent{
		Key:       key,
		NewState:  newState,
		Timestamp: time.Now(),
	}

	// Notify subscribers
	for _, pid := range subscribers {
		context.Send(pid, event)
	}
}

// InMemoryStateStore is an in-memory implementation of StateStore
type InMemoryStateStore struct {
	// states is the state data
	states map[string][]byte
	// checkpoints are the checkpoints
	checkpoints map[string]map[string][]byte
	// mu is a mutex for protecting the state
	mu sync.RWMutex
}

// NewInMemoryStateStore creates a new in-memory state store
func NewInMemoryStateStore() *InMemoryStateStore {
	return &InMemoryStateStore{
		states:      make(map[string][]byte),
		checkpoints: make(map[string]map[string][]byte),
	}
}

// Get gets a state value
func (s *InMemoryStateStore) Get(key string) ([]byte, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.states[key], nil
}

// Put puts a state value
func (s *InMemoryStateStore) Put(key string, value []byte) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.states[key] = value
	return nil
}

// Delete deletes a state value
func (s *InMemoryStateStore) Delete(key string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	delete(s.states, key)
	return nil
}

// CreateCheckpoint creates a checkpoint
func (s *InMemoryStateStore) CreateCheckpoint(id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Create a copy of the current state
	checkpoint := make(map[string][]byte)
	for k, v := range s.states {
		valueCopy := make([]byte, len(v))
		copy(valueCopy, v)
		checkpoint[k] = valueCopy
	}

	// Store the checkpoint
	s.checkpoints[id] = checkpoint
	return nil
}

// RestoreFromCheckpoint restores from a checkpoint
func (s *InMemoryStateStore) RestoreFromCheckpoint(id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	// Check if checkpoint exists
	checkpoint, ok := s.checkpoints[id]
	if !ok {
		return fmt.Errorf("checkpoint not found: %s", id)
	}

	// Restore state from checkpoint
	s.states = make(map[string][]byte)
	for k, v := range checkpoint {
		valueCopy := make([]byte, len(v))
		copy(valueCopy, v)
		s.states[k] = valueCopy
	}

	return nil
}

// ListCheckpoints lists all checkpoints
func (s *InMemoryStateStore) ListCheckpoints() ([]string, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	checkpoints := make([]string, 0, len(s.checkpoints))
	for id := range s.checkpoints {
		checkpoints = append(checkpoints, id)
	}

	return checkpoints, nil
}

// UnsubscribeFromStateChangesCommand is a command to unsubscribe from state changes
type UnsubscribeFromStateChangesCommand struct {
	// Key is the state key to unsubscribe from
	Key string
}

// UnsubscribeFromStateChangesResponse is the response to an UnsubscribeFromStateChangesCommand
type UnsubscribeFromStateChangesResponse struct {
	// Success indicates if the operation was successful
	Success bool
	// Error is the error if the operation failed
	Error error
}

// ListCheckpointsCommand is a command to list checkpoints
type ListCheckpointsCommand struct{}

// ListCheckpointsResponse is the response to a ListCheckpointsCommand
type ListCheckpointsResponse struct {
	// Success indicates if the operation was successful
	Success bool
	// Error is the error if the operation failed
	Error error
	// Checkpoints are the checkpoint IDs
	Checkpoints []string
}

// CheckpointTimerMessage is a message for the checkpoint timer
type CheckpointTimerMessage struct{}
